
import yaml

import os
import psycopg2

from core.database import OpenGaussManager

current_path = os.path.dirname(os.path.realpath(__file__))

"""
# Reference: 
# 1. https://zhuanlan.zhihu.com/p/686030132
# 2. https://docs.opengauss.org/zh/docs/6.0.0/docs/AIFeatureGuide/%E5%8E%9F%E7%94%9FDB4AI%E5%BC%95%E6%93%8E.html
"""
class IrisModel:
    def __init__(self, conn: psycopg2.extensions.connection):
        self.connection = conn
        self.cursor=conn.cursor()
    def create(self):
        cur = self.cursor
        cur.execute("DROP MODEL iris_model;")
        cur.execute('CREATE MODEL iris_model USING logistic_regression FEATURES sepal_length, sepal_width,petal_length TARGET petal_width<1 FROM iris;')
        self.commit()
    def fetch_all(self):
        cur = self.cursor
        cur.execute('SELECT modelname from gs_model_warehouse;')
        results=cur.fetchall()
        print("modelname: %s" %(results))
    def predict(self):
        cur = self.cursor
        cur.execute('SELECT id, PREDICT BY iris_model (FEATURES sepal_length,sepal_width,petal_length) as "PREDICT" FROM iris limit 20;')
        results=cur.fetchall()
        print("predict result: %s" % (results))
    def _load_data(self, iris_data: str):
        cur = self.cursor
        print("load data")
        cur.execute("DROP TABLE IF EXISTS iris;")
        cur.execute("CREATE TABLE iris (id integer, sepal_length float, sepal_width float, petal_length float, petal_width float, species varchar(255));")
        
        #file = open("/home/omm/Iris.csv", "r")
        file = open(iris_data, "r")
        line = file.readline()
        
        line = file.readline()
        while line:
            data = line.split(",")
            id = int(data[0])
            sepal_length = float(data[1])
            sepal_width = float(data[2])
            petal_length = float(data[3])
            petal_width = float(data[4])
            species = data[5].replace("\n", "")
            cmd = "INSERT INTO iris (id, sepal_length, sepal_width, petal_length, petal_width, species) VALUES ({}, {}, {}, {}, {}, '{}');".format(id, sepal_length, sepal_width, petal_length, petal_width, species)
            #print(cmd)
            cur.execute(cmd)
            line = file.readline()
        self.commit()
    def drop(self):
        cur = self.cursor
        cur.execute('DROP MODEL iris_model;')
        self.commit()
    def commit(self):
        self.connection.commit()

# 使用示例
if __name__ == "__main__":
    #print(current_path+"/../data/Iris.csv")
    try:
        pg_manager=OpenGaussManager(current_path+"/config.yaml")
        try:
            conn = pg_manager.connect()
            # 使用连接执行查询
            with conn.cursor() as cursor:
                cursor.execute("SELECT version()")
                print("PostgreSQL版本:", cursor.fetchone())
            print("\n连接成功！")
            model=IrisModel(conn)
            model._load_data(current_path+"/../data/Iris.csv")
            model.create()
            model.fetch_all()
            model.predict()
        except ConnectionError as e:
            print(f"\n连接失败: {e}")
        finally:
            pg_manager.close()
            print("\n连接已关闭")
    except ValueError as e:
        print(f"配置错误: {e}")
    except Exception as e:
        print(f"意外错误: {e}")
    
